#include "CEventTracer.h"

BOOL EscalatePrivilege()
{
    HANDLE hToken = nullptr;
    LUID luid = { 0 };
    TOKEN_PRIVILEGES TokenPrivileges = { 0 };
    BOOL bRet = FALSE;
    DWORD dwRet;

    do
    {
        bRet = ::OpenProcessToken(GetCurrentProcess(), TOKEN_ADJUST_PRIVILEGES, &hToken);
        if (!bRet) break;

        bRet = ::LookupPrivilegeValue(nullptr, SE_DEBUG_NAME, &luid);
        if (!bRet) break;

        TokenPrivileges.PrivilegeCount = 1;
        TokenPrivileges.Privileges[0].Luid = luid;
        TokenPrivileges.Privileges[0].Attributes = SE_PRIVILEGE_ENABLED;
        bRet = ::AdjustTokenPrivileges(hToken, FALSE, &TokenPrivileges, 0, nullptr, nullptr);
        if (!bRet) break;

        dwRet = ::GetLastError();
        if (ERROR_NOT_ALL_ASSIGNED == dwRet)
            bRet = FALSE;
        else
            bRet = TRUE;

    } while (FALSE);

    return bRet;
}

DWORD WINAPI ProcessTraceProc(_In_ LPVOID lpParameter)
{
    FILETIME NowTime;
    ::GetSystemTimeAsFileTime(&NowTime);
    ProcessTrace((PTRACEHANDLE)lpParameter, 1, &NowTime, nullptr);
    return 0;
}

void GetEventName(PEVENT_RECORD EventRecord)
{
    ULONG uBuffSize = 0;
    PTRACE_EVENT_INFO lpTraceEventInfo = nullptr;
    std::wstring strName;

    do
    {
        if (ERROR_INSUFFICIENT_BUFFER != ::TdhGetEventInformation(EventRecord, 0, nullptr, nullptr, &uBuffSize)) break;

        lpTraceEventInfo = (PTRACE_EVENT_INFO)VirtualAlloc(nullptr, uBuffSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
        if (!lpTraceEventInfo) break;

        if (ERROR_SUCCESS != ::TdhGetEventInformation(EventRecord, 0, nullptr, lpTraceEventInfo, &uBuffSize)) break;

        strName = std::wstring((PCWSTR)((PBYTE)lpTraceEventInfo + lpTraceEventInfo->TaskNameOffset)) + L"/" + 
            std::wstring((PCWSTR)((PBYTE)lpTraceEventInfo + lpTraceEventInfo->OpcodeNameOffset));

        wprintf(_T("%s\r\n"), strName.c_str());

    } while (FALSE);

    if (lpTraceEventInfo)
        VirtualFree(lpTraceEventInfo, 0, MEM_RELEASE);
}

void OnEventRecord(IN PEVENT_RECORD EventRecord)
{
    GetEventName(EventRecord);
}

CEventTracer::CEventTracer()
{
    EscalatePrivilege();

    m_EventType = EVENT_FLAG_NONE;
    m_lpEventTraceProperties = nullptr;
    m_hSessionTrace = 0;
    m_hEtwTrace = 0;
    m_hThread = nullptr;
}

CEventTracer::~CEventTracer()
{
    StopEventTrace();
}

void CEventTracer::AddEventType(IN int Type)
{
    m_EventType = (EVENT_FLAGS)Type;
}

ULONG CEventTracer::SetEventFlags()
{
    if (m_EventType != EVENT_FLAG_NONE)
        return TraceSetInformation(m_hSessionTrace, TraceSystemTraceEnableFlagsInfo, &m_EventType, sizeof(m_EventType));

    return 0;
}

BOOL CEventTracer::StartEventTrace(IN LPCTSTR lpInstanceName)
{
    BOOL bSuccess = FALSE;
    LPVOID lpBuff = nullptr;
    size_t stInstanceName;
    size_t stBuffSize;
    PEVENT_TRACE_PROPERTIES lpEventTraceProperties = nullptr;
    ULONG uRet;
    EVENT_TRACE_LOGFILE EventTraceLogFile = { 0 };
    TRACEHANDLE hSessionTrace = 0;
    TRACEHANDLE hEtwTrace = 0;

    do
    {
        if (!lpInstanceName) break;

        stInstanceName = (_tcslen(lpInstanceName) + 1) * sizeof(TCHAR);
        stBuffSize = sizeof(EVENT_TRACE_PROPERTIES) + stInstanceName;
        lpBuff = VirtualAlloc(nullptr, stBuffSize, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE);
        if (!lpBuff) break;
        lpEventTraceProperties = (PEVENT_TRACE_PROPERTIES)lpBuff;

    _start_again:
        RtlZeroMemory(lpEventTraceProperties, stBuffSize);
        lpEventTraceProperties->EnableFlags = EVENT_FLAG_PROCESS;
        lpEventTraceProperties->Wnode.BufferSize = (ULONG)stBuffSize;
        lpEventTraceProperties->Wnode.Guid = EventTracePropertiesGuid;
        lpEventTraceProperties->Wnode.Flags = WNODE_FLAG_TRACED_GUID;
        lpEventTraceProperties->Wnode.ClientContext = 1;
        lpEventTraceProperties->FlushTimer = 1;
        lpEventTraceProperties->LogFileMode = EVENT_TRACE_REAL_TIME_MODE | EVENT_TRACE_USE_LOCAL_SEQUENCE | EVENT_TRACE_SYSTEM_LOGGER_MODE;
        lpEventTraceProperties->LoggerNameOffset = sizeof(EVENT_TRACE_PROPERTIES);
        RtlCopyMemory(lpEventTraceProperties + 1, lpInstanceName, stInstanceName);

        uRet = StartTrace(&hSessionTrace, lpInstanceName, lpEventTraceProperties);
        if (ERROR_SUCCESS != uRet)
        {
            if (ERROR_ALREADY_EXISTS != uRet) break;

            uRet = ControlTrace(hSessionTrace, lpInstanceName, lpEventTraceProperties, EVENT_TRACE_CONTROL_STOP);
            if (uRet != ERROR_SUCCESS) break;

            goto _start_again;
        }

        SetEventFlags();

        EventTraceLogFile.Context = this;
        EventTraceLogFile.LoggerName = (LPTSTR)lpInstanceName;
        EventTraceLogFile.ProcessTraceMode = PROCESS_TRACE_MODE_EVENT_RECORD | PROCESS_TRACE_MODE_REAL_TIME;
        EventTraceLogFile.EventRecordCallback = OnEventRecord;
        hEtwTrace = OpenTrace(&EventTraceLogFile);
        if (INVALID_PROCESSTRACE_HANDLE == hEtwTrace) break;

        m_hEtwTrace = hEtwTrace;

        m_hThread = CreateThread(nullptr, 0, ProcessTraceProc, &m_hEtwTrace, 0, nullptr);
        if (!m_hThread) break;
        SetThreadPriority(m_hThread, THREAD_PRIORITY_HIGHEST);

        m_hSessionTrace = hSessionTrace;
        m_lpEventTraceProperties = lpEventTraceProperties;

        bSuccess = TRUE;

    } while (FALSE);

    if (!bSuccess && hEtwTrace)
    {
        CloseTrace(hEtwTrace);
        m_hEtwTrace = 0;
    }

    if (!bSuccess && hSessionTrace && lpEventTraceProperties)
    {
        ControlTrace(hSessionTrace, (PTCHAR)(lpEventTraceProperties + 1), lpEventTraceProperties, EVENT_TRACE_CONTROL_STOP);
        hSessionTrace = 0;
    }

    if (!bSuccess && lpBuff)
        VirtualFree(lpBuff, 0, MEM_RELEASE);

    return bSuccess;
}

void CEventTracer::StopEventTrace()
{
    if (m_hThread)
    {
        TerminateThread(m_hThread, 0);
        CloseHandle(m_hThread);
        m_hThread = nullptr;
    }

    if (m_hEtwTrace)
    {
        CloseTrace(m_hEtwTrace);
        m_hEtwTrace = 0;
    }

    if (m_hSessionTrace && m_lpEventTraceProperties)
    {
        ControlTrace(m_hSessionTrace, (PTCHAR)(m_lpEventTraceProperties + 1), m_lpEventTraceProperties, EVENT_TRACE_CONTROL_STOP);
        m_hSessionTrace = 0;
    }

    if (m_lpEventTraceProperties)
    {
        VirtualFree(m_lpEventTraceProperties, 0, MEM_RELEASE);
        m_lpEventTraceProperties = nullptr;
    }
}
